from lib.tencent.hunyuan_lite import HunyuanLite
from lib.baidu.ernie_speed import ErnieSpeed
from tools import question


class PromptI:

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "content": (
                    "STRING",
                    {
                        "multiline": True,
                        "default": "",
                    },
                ),
                "passage": ("STRING", {"multiline": True, "default": ""}),
                "engine": (
                    ["ernie_speed_128k", "hunyuan_lite"],
                    {"default": "ernie_speed_128k"},
                ),
            }
        }

    RETURN_TYPES = ("STRING", "STRING")
    RETURN_NAMES = ("result", "question")

    FUNCTION = "handle"

    OUTPUT_NODE = False

    CATEGORY = "fcl/prompt_i"

    history_dict = {}

    def handle(self, content="", passage="", engine=""):

        messages, question = self.assimble(content, passage)
        print(messages)
        if engine == "hunyuan_lite":
            hyl = HunyuanLite()
            ret_msg = hyl.getMessage(hyl.chat(messages=messages))
        else:
            es = ErnieSpeed()
            ret_msg = es.getMessage(es.chat(messages=messages))

        return (ret_msg, question)

    def assimble(self, content, passage):
        question_content = f"""
我有一段文案，然后我要从文案中抽取一段内容来生成一张图片，并且这段内容生成的图片要贴合原文案相应的内容，

图片提示词内容大致为：(
人物描述：[具体的外貌特征，如帅气、阳光等]的[性别]孩，身着[服装款式和颜色]，[发型特征]
动作姿态：正在[具体的动作，如全力奔跑、慢跑等]，[身体部位的动作，如双臂摆动、双腿迈开等]
表情神态：[具体的神情，如专注、坚定、自信等]的表情，[面部细节，如汗水流淌、眉头紧皱等]
环境背景：在[具体的地点，如学校操场、公园跑道等]，[时间和天气，如夕阳西下、阳光明媚等]
其他细节：[其他需要突出的细节，如衣服湿透、鞋带松开等]。
)

以下为一个例子：【

[文案内容]：在空旷的操场上，有一个男孩正全力奔跑着。他身着简单的运动装，身影在夕阳的映照下显得格外矫健。
他的步伐轻盈而有力，每一次脚掌与地面的接触，都仿佛在释放着内心积压的能量。风掠过他的脸颊，吹乱了他的头发，却吹不散他专注的神情。
旁人的目光并未影响到他，他沉浸在自己的节奏中。呼吸略显急促，但那坚定的眼神表明他不会轻易停下。
他的双臂有规律地摆动着，像是在指挥着一场只属于他自己的交响乐。操场上的标线一格一格被他甩在身后，仿佛是他征服的领土。
这个男孩，或许是在为即将到来的比赛做准备，或许只是在享受奔跑带来的自由。但无论原因如何，他此刻的坚持和执着，都让这片操场充满了生机与活力。
汗水湿透了他的衣衫，可他浑然不觉，只是不断地向前、向前。在他的奔跑中，仿佛能看到对挑战的无畏，对目标的渴望。
夕阳渐渐西沉，将他的身影拉得更长，而他的脚步，依旧不停歇，仿佛要一直跑向那未知的远方。

[抽取的内容]：汗水湿透了他的衣衫，可他浑然不觉，只是不断地向前、向前。在他的奔跑中。

[图片提示词]：
人物描述：阳光帅气的男孩，身着白色运动 T 恤和黑色运动短裤，短发利落
动作姿态：正在全力奔跑，双臂有节奏地摆动
表情神态：神情专注，汗水顺着脸颊流淌
环境背景：在学校的操场上，夕阳的余晖映照
其他细节：衣服被汗水湿透

】。

[文案内容]:

{content}

[抽取内容]：
{passage}

请返回[图片提示词]。



        """
        message = {
            "messages": [
                {
                    "role": "user",
                    "content": question_content,
                },
            ]
        }
        return message, question_content


NODE_CLASS_MAPPINGS = {"FCL_PROMPT_I": PromptI}

NODE_DISPLAY_NAME_MAPPINGS = {"FCL_PROMPT_I": "fcl dev prompt 1"}
